package com.inspur.edp.ai.flow.core.engine;

import com.inspur.edp.ai.flow.engine.*;
import com.inspur.edp.ai.flow.exception.FlowException;
import com.inspur.edp.ai.flow.metadata.AbsNode;
import com.inspur.edp.ai.flow.metadata.Edge;
import com.inspur.edp.common.expr.evaluator.ExpressEvaluator;
import lombok.Data;

import java.util.HashMap;
import java.util.Map;

/**
 * @author lizhaorui
 * @date 2025/9/11
 * @description
 */

@Data
public class EngineContext {

    /**
     * 编排中所有流程，主流程、子流程
     */

    private Map<String, Flow> flows;

    /**
     * 流程上下文
     */
    private FlowContext flowContext;


    /**
     * 数据总线
     */
    private DataBus dataBus;


    private Map<String, Object> returnValues;

    /**
     * 节点编译结果
     * Key ：nodeId
     */
    Map<String, Map<String, ExpressEvaluator>> inputParamEvaluators;


    /**
     * Key为 flowId value.key nodeId
     */
    private Map<String, Map<String, NodeEdgeArriveRecord>> flowNodeEdgeArriveRecords;


    /**
     * 结束监听 Key：flowId
     */
    private Map<String, FlowFinRecord> flowFinRecords;


    private boolean returned;


    public EngineContext() {
        this.flows = new HashMap<>();
        this.flowNodeEdgeArriveRecords = new HashMap<>();
        this.flowFinRecords = new HashMap<>();
        this.returnValues = new HashMap<>();
    }

    public NodeEdgeArriveRecord addArriveEdge(String flowId, Edge edge, boolean isEmptyPipe) {
        String targetNodeId = edge.getTargetNodeId();
        NodeEdgeArriveRecord edgeArriveRecord = null;
        Map<String, NodeEdgeArriveRecord> nodeEdgeArriveRecords = this.flowNodeEdgeArriveRecords.get(flowId);
        if (!nodeEdgeArriveRecords.containsKey(targetNodeId)) {
            edgeArriveRecord = new NodeEdgeArriveRecord();
            nodeEdgeArriveRecords.put(targetNodeId, edgeArriveRecord);
        } else {
            edgeArriveRecord = nodeEdgeArriveRecords.get(targetNodeId);
        }
        edgeArriveRecord.addEdge(edge, isEmptyPipe);
        return edgeArriveRecord;
    }

    public FlowFinRecord addFinNode(String flowId, String endNodeId) {
        FlowFinRecord endNodeRecord = this.flowFinRecords.get(flowId);
        endNodeRecord.getFinishedEndNodeIds().add(endNodeId);
        return endNodeRecord;
    }

    public Map<String, Object> getNodeInputParamValues(AbsNode node) {
        Map<String, ExpressEvaluator> paramEvaluators = this.inputParamEvaluators.get(node.getId());
        if (paramEvaluators == null) {
            return null;
        }
        Map<String, Object> retMap = new HashMap<>();
        for (Map.Entry<String, ExpressEvaluator> entry : paramEvaluators.entrySet()) {
            Object paramValue = entry.getValue().evaluate(this.getDataBus().getDataMap());
            retMap.put(entry.getKey(), paramValue);
        }
        //把自身的数据也放入
        if (node.isInputSelf()) {
            retMap.put(node.getCode(), this.getDataBus().get(node.getCode()));
        }
        return retMap;
    }

    public void addFlow(Flow flow) {
        this.flows.put(flow.getId(), flow);
        FlowFinRecord flowFinRecord = new FlowFinRecord();
        flowFinRecord.setEndNodeIds(flow.getEndNodeIds());
        this.flowFinRecords.put(flow.getId(), flowFinRecord);
        this.flowNodeEdgeArriveRecords.put(flow.getId(), new HashMap<>());
    }

    public void switchToParentFlow(Flow flow) {
        this.deleteFlow(flow);
        String parentFlowId = flow.getParentId();
        Flow parentFlow = this.getFlow(parentFlowId);
        this.getFlowContext().setFlow(parentFlow);
    }

    private void deleteFlow(Flow flow) {
        String flowId = flow.getId();
        // 删除流程节点 DataBus数据
        if (flow.getNodes() != null) {
            for (AbsNode node : flow.getNodes()) {
                this.dataBus.remove(node.getCode());
            }
        }
        this.flows.remove(flowId);
        this.flowFinRecords.remove(flowId);
        this.flowNodeEdgeArriveRecords.remove(flowId);
    }

    private Flow getFlow(String flowId) {
        Flow flow = this.flows.get(flowId);
        if (flow == null) {
            throw new FlowException(String.format("流程ID：%s 不存在。", flowId));
        }
        return flow;
    }

}